import logging
from functools import partial

from datasets.patch_pair_dataset import PatchPairDataset
from datasets.patch_clipper import PatchClipper

logger = logging.getLogger("match")

class PairMatchDataset(PatchPairDataset, PatchClipper):


    def __init__(self, config=None, *, tools=None, name="", triplet=None):
        super(PairMatchDataset, self).__init__(config, super=True)
        clipper_setting = config.clipper.toDict() if "clipper" in dir(config) else None
        if not clipper_setting is None: 
            super(PatchPairDataset, self).__init__(**clipper_setting)
        else:
            self.crop = lambda x, y : (x, y)
            
        #-- dataset
        self.dt = tools(config.base_dir, config.list_file)
        self.pl = self.dt.get_pair_list()

        self.triplet = triplet

        logger.info(f"{name} Pair Match Dataset created with {self.__len__()} pairs ---- triplet is {self.triplet}")


    def __getitem__(self, index):
        """
        ((sar, opt))
        """
        (sar_tensor, opt_tensor), _ = super().__getitem__(index)

        sar_t, opt_t = self.crop(sar_tensor, opt_tensor)

        ret = (sar_t, opt_t) if self.triplet is None else self.triplet(sar_t, opt_t)

        return ret

